import cv2
import numpy as np
import yaml
from collections import defaultdict
from discret_image import active_size_map_discretization


MODES = ["pgm_to_png","png_red_to_pgm", "node_edge_mode", "auto", "visulize"]
mode = MODES[3]

pgm_path = "../maps/outdoor_final.pgm"
config_path = "../maps/outdoor_final.yaml"
png_path = pgm_path.replace(".pgm", ".png")
graph_path = pgm_path.replace(".pgm", "_graph.yaml")
vis_png_path = pgm_path.replace(".pgm", "_vis.png")
node_radius = 200

def turn_pgm_to_png(pgm_path, png_path):
    img = cv2.imread(pgm_path, -1)
    if img is None:
        raise ValueError(f"无法读取PGM文件: {pgm_path}")

    if len(img.shape) == 2:
        img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    else:
        img_rgb = img[:, :, :3].copy()  # 避免内存重叠问题

    cv2.imwrite(png_path, img_rgb)


def turn_png_red_to_pgm(png_path, pgm_path):
    img = cv2.imread(png_path)
    if img is None:
        raise ValueError(f"无法读取PNG文件: {png_path}")
    red_low = np.array([0, 0, 200])
    red_high = np.array([20, 20, 255])
    mask = cv2.inRange(img, red_low, red_high)
    red_pixels = cv2.bitwise_and(img, img, mask=mask)
    red_gray = cv2.cvtColor(red_pixels, cv2.COLOR_BGR2GRAY)
    red_gray[red_gray > 0] = 255  # 将非零值设为255
    cv2.imwrite(pgm_path, red_gray)
    
def draw_graph_on_image(img, graph):
    nodes = np.array(graph["nodes"], dtype=np.int32)
    edges = graph["edges"]

    img_with_graph = cv2.cvtColor(
        cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR
    )

    for idx in range(len(nodes)):
        cv2.circle(img_with_graph, tuple(nodes[idx][:2]), 1, (0, 255, 0), -1)

    for e in edges:
        cv2.line(
            img_with_graph,
            tuple(nodes[e[0]][:2]),
            tuple(nodes[e[1]][:2]),
            (255, 0, 255),
            1,
        )

    cv2.imshow("Graph",img_with_graph)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    return img_with_graph

def visualize_graph(img_path, graph_path, config_path):
    img = cv2.imread(img_path)
    with open(graph_path, "r") as f:
        graph = yaml.safe_load(f)
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    resolution = config["resolution"]
    x, y, yaw = config["origin"]
    yaw = np.deg2rad(yaw)
    R = np.array([[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]])
    T = np.array([x, y])
    
    # 世界坐标 → 图像坐标（带Y轴翻转）
    height, width = img.shape[:2]
    transformed_nodes = []
    for node in graph["nodes"]:
        world = np.array(node[:2])
        img_coords = R.T @ ((world - T) / resolution)
        x_img, y_img = img_coords
        y_img = height - y_img  # Y轴翻转
        transformed_nodes.append([x_img, y_img])
    graph["nodes"] = transformed_nodes
    draw_graph_on_image(img, graph)

def save_graph_to_yaml(graph, config_path, graph_path):
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    resolution = config["resolution"]
    x, y, yaw = config["origin"]
    yaw = np.deg2rad(yaw)
    R = np.array([[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]])
    T = np.array([x, y])
    
    # 图像坐标 → 世界坐标（带Y轴翻转）
    height, width = cv2.imread(png_path).shape[:2]
    transformed_nodes = []
    for node in graph["nodes"]:
        x_img, y_img = node[:2]
        y_img = height - y_img  # Y轴翻转
        world_coords = R @ (np.array([x_img, y_img]) * resolution) + T
        transformed_nodes.append(world_coords.tolist())
    graph["nodes"] = [[x, y, 0] for x, y in transformed_nodes]
    
    with open(graph_path, "w") as f:
        yaml.dump(graph, f, default_flow_style=False)

def get_graph_from_image(png_path, config_path, graph_path):
    """
    Extracts a graph from an image where:
    - Nodes are red dots [0, 0, 255] in BGR
    - Edges are blue lines [255, 0, 0] in BGR
    - Converts the graph to a real-world coordinate system using config.
    """
    # Load image
    img = cv2.imread(png_path)

    # --- Step 1: Node Detection (Red Dots) ---
    lower_red = np.array([0, 0, 200])
    upper_red = np.array([0, 0, 255])
    mask_red = cv2.inRange(img, lower_red, upper_red)

    # Find contours of red regions
    contours_red, _ = cv2.findContours(
        mask_red, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    nodes = []
    for cnt in contours_red:
        M = cv2.moments(cnt)
        if M["m00"] != 0:
            x = int(M["m10"] / M["m00"])
            y = int(M["m01"] / M["m00"])
            nodes.append([x, y])

    # --- Step 2: Edge Detection (Blue Lines) ---
    lower_blue = np.array([200, 0, 0])
    upper_blue = np.array([255, 0, 0])
    mask_blue = cv2.inRange(img, lower_blue, upper_blue)

    # Morphological closing to connect broken lines
    kernel = np.ones((3, 3), np.uint8)
    mask_blue = cv2.morphologyEx(mask_blue, cv2.MORPH_CLOSE, kernel)

    # Find all blue pixels
    points = cv2.findNonZero(mask_blue)
    if points is None or len(points) < 10:
        edges = []
    else:
        points = points.reshape(-1, 2).astype(np.float32)

        # 分割连通区域
        num_labels, labels = cv2.connectedComponents(mask_blue)

        # 将点按照连通区域标签分组
        label_to_points = defaultdict(list)
        for point in points:
            x, y = int(point[0]), int(point[1])
            label = labels[y, x]
            label_to_points[label].append([x, y])

        edges = []

        # 对每个连通区域单独处理
        for label in range(1, num_labels):  # 跳过背景（label=0）
            cluster_points = np.array(label_to_points[label], dtype=np.float32)

            if len(cluster_points) < 2:
                continue

            # 使用PCA拟合直线方向
            mean, eig, _ = cv2.PCACompute2(cluster_points, mean=np.array([]))
            direction = eig[0]

            # 投影到主方向上
            vectors = cluster_points - mean[0]
            projections = np.dot(vectors, direction)

            min_idx = np.argmin(projections)
            max_idx = np.argmax(projections)
            pt1 = tuple(cluster_points[min_idx].astype(int))
            pt2 = tuple(cluster_points[max_idx].astype(int))

            # 寻找最近的两个节点
            min_dist1 = float("inf")
            min_idx1 = -1
            min_dist2 = float("inf")
            min_idx2 = -1

            for idx, node in enumerate(nodes):
                dist1 = np.linalg.norm(np.array(pt1) - np.array(node))
                dist2 = np.linalg.norm(np.array(pt2) - np.array(node))

                if dist1 < min_dist1:
                    min_dist1 = dist1
                    min_idx1 = idx
                if dist2 < min_dist2:
                    min_dist2 = dist2
                    min_idx2 = idx

            edges.append([min_idx1, min_idx2])

    # --- Step 3: Build Graph Structure ---
    graph = {
        "nodes": nodes,
        "edges": edges,
    }
    draw_graph_on_image(img, graph)

    # 保存
    save_graph_to_yaml(graph, config_path, graph_path)
    return graph

def get_graph_by_discretization(img_path, config_path, graph_path, root=[0, 0]):
    """
    通过图像离散化生成图结构，包含图像预处理步骤
    """
    # 1. 读取图像
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    
    # 2. 图像预处理
    print("[预处理] 开始图像降噪和增强...")
    img_denoised = cv2.medianBlur(img, 3)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    img_closed = cv2.morphologyEx(img_denoised, cv2.MORPH_CLOSE, kernel, iterations=1)
    img_smooth = cv2.GaussianBlur(img_closed, (3, 3), 0)
    _, binary_img = cv2.threshold(img_smooth, 150, 255, cv2.THRESH_BINARY)
    
    # 3. 读取配置参数
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    resolution = config["resolution"]
    x, y, yaw = config["origin"]
    # ROS map_saver: yaw为弧度，像素坐标原点在左下角，y轴向上
    # 图像坐标系：左上角为(0,0)，x向右，y向下
    height, width = img.shape[:2]
    # 世界坐标转像素坐标
    root_world = np.array(root[:2], dtype=np.float32)
    # 先减去origin平移，再旋转
    R = np.array([[np.cos(yaw), -np.sin(yaw)],
                  [np.sin(yaw),  np.cos(yaw)]], dtype=np.float32)
    T = np.array([x, y], dtype=np.float32)
    # 世界坐标转到地图坐标
    map_xy = (R.T @ (root_world - T)) / resolution
    # 地图坐标转到像素坐标（y轴翻转）
    root_pixel = np.array([map_xy[0], height - 1 - map_xy[1]])
    root_pixel = np.round(root_pixel).astype(int)
    root_pixel[0] = np.clip(root_pixel[0], 0, img.shape[1] - 1)
    root_pixel[1] = np.clip(root_pixel[1], 0, img.shape[0] - 1)
    root_pixel = tuple(root_pixel)

    # 5. 使用漫水填充标记可行区域
    h, w = img.shape
    mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
    flooded = binary_img.copy()
    flood_color = 100
    lo_diff, hi_diff = 10, 10
    flags = 4 | (lo_diff << 8) | cv2.FLOODFILL_FIXED_RANGE
    cv2.floodFill(flooded, mask, root_pixel, flood_color, lo_diff, hi_diff, flags)

    # 6. 调用离散化函数
    runnable_color = flood_color
    result_img, graph = active_size_map_discretization(flooded, runnable_color=runnable_color, node_radius=node_radius)

    # 7. 找到最大连通子图，并覆盖原图结构
    # 构建邻接表表示图
    adj_list = [[] for _ in range(len(graph["nodes"]))]
    for edge in graph["edges"]:
        u, v = edge
        adj_list[u].append(v)
        adj_list[v].append(u)
    
    # BFS遍历查找连通分量
    visited = [False] * len(graph["nodes"])
    max_component = []
    
    for i in range(len(graph["nodes"])):
        if not visited[i]:
            component = []
            queue = [i]
            visited[i] = True
            
            while queue:
                node = queue.pop(0)
                component.append(node)
                for neighbor in adj_list[node]:
                    if not visited[neighbor]:
                        visited[neighbor] = True
                        queue.append(neighbor)
            
            # 更新最大连通分量
            if len(component) > len(max_component):
                max_component = component
    
    # 创建节点映射表（旧索引->新索引）
    node_map = {}
    new_nodes = []
    for idx, old_idx in enumerate(max_component):
        node_map[old_idx] = idx
        new_nodes.append(graph["nodes"][old_idx])
    
    # 创建新边列表（只保留最大连通分量内的边）
    new_edges = []
    for edge in graph["edges"]:
        u, v = edge
        if u in node_map and v in node_map:
            new_edges.append([node_map[u], node_map[v]])
    
    # 更新图结构
    graph["nodes"] = new_nodes
    graph["edges"] = new_edges
    
    # 8. 可视化结果
    cv2.imwrite(vis_png_path, result_img)

    # 9. 保存图结构
    save_graph_to_yaml(graph, config_path, graph_path)
    return graph


if __name__ == "__main__":
    print("Running mode:", mode)
    if mode == MODES[0]:
        turn_pgm_to_png(pgm_path, png_path)
    elif mode == MODES[1]:
        turn_png_red_to_pgm(png_path, pgm_path)
    elif mode == MODES[2]:
        get_graph_from_image(png_path, config_path, graph_path)
    elif mode == MODES[3]:
        get_graph_by_discretization(png_path, config_path, graph_path)
    else:
        visualize_graph(png_path, graph_path, config_path)